import numpy as np

# 定义KMeans类
class KMeans:
    def __init__(self, k=2, max_iters=100, tolerance=1e-4):
        self.k = k  # 聚类数
        self.max_iters = max_iters  # 最大迭代次数
        self.tolerance = tolerance  # 容忍误差
        self.centroids = None  # 聚类中心

    def fit(self, X):
        # 初始化聚类中心
        np.random.seed(42)
        random_idx = np.random.choice(len(X), size=self.k, replace=False)
        self.centroids = X[random_idx]

        for i in range(self.max_iters):
            # 步骤1: 为每个点分配标签
            labels = self._assign_labels(X)

            # 步骤2: 计算新的聚类中心
            new_centroids = self._compute_centroids(X, labels)

            # 步骤3: 检查收敛情况
            centroid_shift = np.sum((new_centroids - self.centroids) ** 2)
            if centroid_shift <= self.tolerance:
                print(f"Converged in {i+1} iterations.")
                break

            self.centroids = new_centroids

    def _assign_labels(self, X):
        """为每个数据点分配最近的聚类中心"""
        labels = np.argmin(np.linalg.norm(X[:, np.newaxis] - self.centroids, axis=2), axis=1)
        return labels

    def _compute_centroids(self, X, labels):
        """根据数据点的标签重新计算聚类中心"""
        centroids = np.array([X[labels == i].mean(axis=0) for i in range(self.k)])
        return centroids

    def predict(self, X):
        """给定数据点，预测其所属的簇"""
        return self._assign_labels(X)

    def get_centroids(self):
        """返回聚类中心"""
        return self.centroids


# 测试代码
if __name__ == "__main__":
    # 随机生成一些数据点
    X = np.random.randn(100, 2)  # 100个二维数据点
    kmeans = KMeans(k=3)
    kmeans.fit(X)

    # 输出聚类中心
    print("聚类中心：")
    print(kmeans.get_centroids())

    # 测试预测
    labels = kmeans.predict(X)
    print("预测的标签：")
    print(labels)